{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# demo conversion\n",
    "import torch\n",
    "import pickle\n",
    "import numpy as np\n",
    "from hparams import hparams\n",
    "from utils import pad_seq_to_2\n",
    "from utils import quantize_f0_numpy\n",
    "from model import Generator_3 as Generator\n",
    "from model import Generator_6 as F0_Converter\n",
    "\n",
    "\n",
    "device = 'cuda:0'\n",
    "G = Generator(hparams).eval().to(device)\n",
    "g_checkpoint = torch.load('assets/660000-G.ckpt', map_location=lambda storage, loc: storage)\n",
    "G.load_state_dict(g_checkpoint['model'])\n",
    "\n",
    "P = F0_Converter(hparams).eval().to(device)\n",
    "p_checkpoint = torch.load('assets/640000-P.ckpt', map_location=lambda storage, loc: storage)\n",
    "P.load_state_dict(p_checkpoint['model'])\n",
    "\n",
    "\n",
    "metadata = pickle.load(open('assets/demo.pkl', \"rb\"))\n",
    "\n",
    "\n",
    "sbmt_i = metadata[0]\n",
    "emb_org = torch.from_numpy(sbmt_i[1]).to(device)\n",
    "x_org, f0_org, len_org, uid_org = sbmt_i[2]        \n",
    "uttr_org_pad, len_org_pad = pad_seq_to_2(x_org[np.newaxis,:,:], 192)\n",
    "uttr_org_pad = torch.from_numpy(uttr_org_pad).to(device)\n",
    "f0_org_pad = np.pad(f0_org, (0, 192-len_org), 'constant', constant_values=(0, 0))\n",
    "f0_org_quantized = quantize_f0_numpy(f0_org_pad)[0]\n",
    "f0_org_onehot = f0_org_quantized[np.newaxis, :, :]\n",
    "f0_org_onehot = torch.from_numpy(f0_org_onehot).to(device)\n",
    "uttr_f0_org = torch.cat((uttr_org_pad, f0_org_onehot), dim=-1)\n",
    "\n",
    "sbmt_j = metadata[1]\n",
    "emb_trg = torch.from_numpy(sbmt_j[1]).to(device)\n",
    "x_trg, f0_trg, len_trg, uid_trg = sbmt_j[2]        \n",
    "uttr_trg_pad, len_trg_pad = pad_seq_to_2(x_trg[np.newaxis,:,:], 192)\n",
    "uttr_trg_pad = torch.from_numpy(uttr_trg_pad).to(device)\n",
    "f0_trg_pad = np.pad(f0_trg, (0, 192-len_trg), 'constant', constant_values=(0, 0))\n",
    "f0_trg_quantized = quantize_f0_numpy(f0_trg_pad)[0]\n",
    "f0_trg_onehot = f0_trg_quantized[np.newaxis, :, :]\n",
    "f0_trg_onehot = torch.from_numpy(f0_trg_onehot).to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    f0_pred = P(uttr_org_pad, f0_trg_onehot)[0]\n",
    "    f0_pred_quantized = f0_pred.argmax(dim=-1).squeeze(0)\n",
    "    f0_con_onehot = torch.zeros((1, 192, 257), device=device)\n",
    "    f0_con_onehot[0, torch.arange(192), f0_pred_quantized] = 1\n",
    "uttr_f0_trg = torch.cat((uttr_org_pad, f0_con_onehot), dim=-1)    \n",
    "\n",
    "\n",
    "conditions = ['R', 'F', 'U', 'RF', 'RU', 'FU', 'RFU']\n",
    "spect_vc = []\n",
    "with torch.no_grad():\n",
    "    for condition in conditions:\n",
    "        if condition == 'R':\n",
    "            x_identic_val = G(uttr_f0_org, uttr_trg_pad, emb_org)\n",
    "        if condition == 'F':\n",
    "            x_identic_val = G(uttr_f0_trg, uttr_org_pad, emb_org)\n",
    "        if condition == 'U':\n",
    "            x_identic_val = G(uttr_f0_org, uttr_org_pad, emb_trg)\n",
    "        if condition == 'RF':\n",
    "            x_identic_val = G(uttr_f0_trg, uttr_trg_pad, emb_org)\n",
    "        if condition == 'RU':\n",
    "            x_identic_val = G(uttr_f0_org, uttr_trg_pad, emb_trg)\n",
    "        if condition == 'FU':\n",
    "            x_identic_val = G(uttr_f0_trg, uttr_org_pad, emb_trg)\n",
    "        if condition == 'RFU':\n",
    "            x_identic_val = G(uttr_f0_trg, uttr_trg_pad, emb_trg)\n",
    "            \n",
    "        if 'R' in condition:\n",
    "            uttr_trg = x_identic_val[0, :len_trg, :].cpu().numpy()\n",
    "        else:\n",
    "            uttr_trg = x_identic_val[0, :len_org, :].cpu().numpy()\n",
    "                \n",
    "        spect_vc.append( ('{}_{}_{}_{}'.format(sbmt_i[0], sbmt_j[0], uid_org, condition), uttr_trg ) )       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# spectrogram to waveform\n",
    "import torch\n",
    "import librosa\n",
    "import pickle\n",
    "import os\n",
    "from synthesis import build_model\n",
    "from synthesis import wavegen\n",
    "\n",
    "if not os.path.exists('results'):\n",
    "    os.makedirs('results')\n",
    "\n",
    "model = build_model().to(device)\n",
    "checkpoint = torch.load(\"assets/checkpoint_step001000000_ema.pth\")\n",
    "model.load_state_dict(checkpoint[\"state_dict\"])\n",
    "\n",
    "for spect in spect_vc:\n",
    "    name = spect[0]\n",
    "    c = spect[1]\n",
    "    print(name)\n",
    "    waveform = wavegen(model, c=c)   \n",
    "    librosa.output.write_wav('results/'+name+'.wav', waveform, sr=16000)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
